{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup and Installation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/drive/1vysj4eGqYt4sBdMp1BDI5Kxp4jtQbBxZ?usp=sharing\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "\n",
    "# Cookbook LlamaIndex Integration by Maxim AI (Instrumentation Module)\n",
    "\n",
    "This is a simple cookbook that demonstrates how to use the [LlamaIndex Maxim integration](https://www.getmaxim.ai/docs/sdk/python/integrations/llamaindex/llamaindex) using the [instrumentation module](https://docs.llamaindex.ai/en/stable/module_guides/observability/instrumentation/) by LlamaIndex (available in llama-index v0.10.20 and later)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img \n",
    "  src=\"https://cdn.getmaxim.ai/public/images/llamaindex.gif\" \n",
    "  alt=\"LlamaIndex demo gif\"\n",
    "  style=\"max-width: 100%; height: auto; border-radius: 10px;\"\n",
    "/>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install required packages\n",
    "\n",
    "# pip install llama-index\n",
    "# pip install llama-index-llms-openai\n",
    "# pip install llama-index-embeddings-openai\n",
    "# pip install llama-index-tools-wikipedia\n",
    "# pip install llama-index-tools-requests\n",
    "# pip install maxim-py\n",
    "# pip install python-dotenv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "# Load environment variables from .env file\n",
    "load_dotenv()\n",
    "\n",
    "# Get environment variables\n",
    "MAXIM_API_KEY = os.getenv(\"MAXIM_API_KEY\")\n",
    "MAXIM_LOG_REPO_ID = os.getenv(\"MAXIM_LOG_REPO_ID\")\n",
    "OPENAI_API_KEY = os.getenv(\"OPENAI_API_KEY\")\n",
    "\n",
    "# Verify required environment variables are set\n",
    "if not MAXIM_API_KEY:\n",
    "    raise ValueError(\"MAXIM_API_KEY environment variable is required\")\n",
    "if not MAXIM_LOG_REPO_ID:\n",
    "    raise ValueError(\"MAXIM_LOG_REPO_ID environment variable is required\")\n",
    "if not OPENAI_API_KEY:\n",
    "    raise ValueError(\"OPENAI_API_KEY environment variable is required\")\n",
    "\n",
    "print(\"✅ Environment variables loaded successfully\")\n",
    "print(\n",
    "    f\"MAXIM_API_KEY: {'*' * (len(MAXIM_API_KEY) - 4) + MAXIM_API_KEY[-4:] if MAXIM_API_KEY else 'Not set'}\"\n",
    ")\n",
    "print(f\"MAXIM_LOG_REPO_ID: {MAXIM_LOG_REPO_ID}\")\n",
    "print(\n",
    "    f\"OPENAI_API_KEY: {'*' * (len(OPENAI_API_KEY) - 4) + OPENAI_API_KEY[-4:] if OPENAI_API_KEY else 'Not set'}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Maxim Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "from maxim import Config, Maxim\n",
    "from maxim.logger import LoggerConfig\n",
    "from maxim.logger.llamaindex import instrument_llamaindex\n",
    "\n",
    "# Initialize Maxim logger\n",
    "maxim = Maxim(Config(api_key=os.getenv(\"MAXIM_API_KEY\")))\n",
    "logger = maxim.logger(LoggerConfig(id=os.getenv(\"MAXIM_LOG_REPO_ID\")))\n",
    "\n",
    "# Instrument LlamaIndex with Maxim observability\n",
    "# Set debug=True to see detailed logs during development\n",
    "instrument_llamaindex(logger)\n",
    "\n",
    "print(\"✅ Maxim instrumentation enabled for LlamaIndex\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simple FunctionAgent with Observability"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.agent import FunctionAgent\n",
    "from llama_index.core.tools import FunctionTool\n",
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "\n",
    "# Define simple calculator tools\n",
    "def add_numbers(a: float, b: float) -> float:\n",
    "    \"\"\"Add two numbers together.\"\"\"\n",
    "    return a + b\n",
    "\n",
    "\n",
    "def multiply_numbers(a: float, b: float) -> float:\n",
    "    \"\"\"Multiply two numbers together.\"\"\"\n",
    "    return a * b\n",
    "\n",
    "\n",
    "def divide_numbers(a: float, b: float) -> float:\n",
    "    \"\"\"Divide first number by second number.\"\"\"\n",
    "    if b == 0:\n",
    "        raise ValueError(\"Cannot divide by zero\")\n",
    "    return a / b\n",
    "\n",
    "\n",
    "# Create function tools\n",
    "add_tool = FunctionTool.from_defaults(fn=add_numbers)\n",
    "multiply_tool = FunctionTool.from_defaults(fn=multiply_numbers)\n",
    "divide_tool = FunctionTool.from_defaults(fn=divide_numbers)\n",
    "\n",
    "# Initialize LLM\n",
    "llm = OpenAI(model=\"gpt-4o-mini\", temperature=0)\n",
    "\n",
    "# Create FunctionAgent\n",
    "agent = FunctionAgent(\n",
    "    tools=[add_tool, multiply_tool, divide_tool],\n",
    "    llm=llm,\n",
    "    verbose=True,\n",
    "    system_prompt=\"\"\"You are a helpful calculator assistant.\n",
    "    Use the provided tools to perform mathematical calculations.\n",
    "    Always explain your reasoning step by step.\"\"\",\n",
    ")\n",
    "\n",
    "# Test the agent with a complex calculation\n",
    "import asyncio\n",
    "\n",
    "\n",
    "async def test_function_agent():\n",
    "    print(\"🔍 Testing FunctionAgent with Maxim observability...\")\n",
    "\n",
    "    query = \"What is (15 + 25) multiplied by 2, then divided by 8?\"\n",
    "\n",
    "    print(f\"\\n📝 Query: {query}\")\n",
    "\n",
    "    # This will be automatically logged by Maxim instrumentation\n",
    "    # FunctionAgent.run() is async, so we need to await it\n",
    "    response = await agent.run(query)\n",
    "\n",
    "    print(f\"\\n🤖 Response: {response}\")\n",
    "    print(\"\\n✅ Check your Maxim dashboard for detailed trace information!\")\n",
    "\n",
    "\n",
    "# Run the async function\n",
    "await test_function_agent()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multi Modal Requests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.agent.workflow import FunctionAgent\n",
    "from llama_index.core.llms import ChatMessage, ImageBlock, TextBlock\n",
    "from llama_index.llms.openai import OpenAI\n",
    "import requests\n",
    "from PIL import Image\n",
    "import io\n",
    "import base64\n",
    "\n",
    "\n",
    "# Tool for image analysis\n",
    "def describe_image_content(description: str) -> str:\n",
    "    \"\"\"Analyze and describe what's in an image based on the model's vision.\"\"\"\n",
    "    return f\"Image analysis complete: {description}\"\n",
    "\n",
    "\n",
    "# Math tools for the agent\n",
    "def add(a: int, b: int) -> int:\n",
    "    \"\"\"Add two numbers together.\"\"\"\n",
    "    return a + b\n",
    "\n",
    "\n",
    "def multiply(a: int, b: int) -> int:\n",
    "    \"\"\"Multiply two numbers together.\"\"\"\n",
    "    return a * b\n",
    "\n",
    "\n",
    "# Create multi-modal agent with vision-capable model\n",
    "multimodal_llm = OpenAI(model=\"gpt-4o-mini\")  # Vision-capable model\n",
    "\n",
    "multimodal_agent = FunctionAgent(\n",
    "    tools=[add, multiply, describe_image_content],\n",
    "    llm=multimodal_llm,\n",
    "    system_prompt=\"You are a helpful assistant that can analyze images and perform calculations.\",\n",
    ")\n",
    "\n",
    "\n",
    "async def test_multimodal_agent():\n",
    "    print(\"🔍 Testing Multi-Modal Agent with Maxim observability...\")\n",
    "\n",
    "    # Create a simple test image (you can replace this with an actual image path)\n",
    "    # For demo purposes, we'll create a simple mathematical equation image\n",
    "    try:\n",
    "        # You can replace this with a real image path if available\n",
    "        # For now, we'll use text-based interaction\n",
    "        # text_query = \"Calculate 15 + 25 and then multiply the result by 3\"\n",
    "\n",
    "        # response = await multimodal_agent.run(text_query)\n",
    "        # print(f\"\\n🤖 Text Response: {response}\")\n",
    "\n",
    "        # If you have an image, you can use this pattern:\n",
    "        msg = ChatMessage(\n",
    "            role=\"user\",\n",
    "            blocks=[\n",
    "                TextBlock(\n",
    "                    text=\"What do you see in this image? If there are numbers, perform calculations.\"\n",
    "                ),\n",
    "                ImageBlock(\n",
    "                    url=\"https://www.shutterstock.com/image-photo/simple-mathematical-equation-260nw-350386472.jpg\"\n",
    "                ),  # Replace with actual image path\n",
    "            ],\n",
    "        )\n",
    "        response = await multimodal_agent.run(msg)\n",
    "\n",
    "    except Exception as e:\n",
    "        print(\n",
    "            f\"Note: Multi-modal features require actual image files. Error: {e}\"\n",
    "        )\n",
    "        print(\n",
    "            \"The agent structure is set up correctly for when you have images to process!\"\n",
    "        )\n",
    "\n",
    "    print(\"\\n✅ Check Maxim dashboard for multi-modal agent traces!\")\n",
    "\n",
    "\n",
    "# Run the test\n",
    "await test_multimodal_agent()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multiple Agents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.agent.workflow import AgentWorkflow, FunctionAgent\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.core.tools import FunctionTool  # Import FunctionTool\n",
    "\n",
    "\n",
    "# Research agent tools\n",
    "def research_topic(topic: str) -> str:\n",
    "    \"\"\"Research a given topic and return key findings.\"\"\"\n",
    "    # Mock research results - in production, this would call real APIs\n",
    "    research_data = {\n",
    "        \"climate change\": \"Climate change refers to long-term shifts in global temperatures and weather patterns, primarily caused by human activities since the 1800s.\",\n",
    "        \"renewable energy\": \"Renewable energy comes from sources that are naturally replenishing like solar, wind, hydro, and geothermal power.\",\n",
    "        \"artificial intelligence\": \"AI involves creating computer systems that can perform tasks typically requiring human intelligence.\",\n",
    "        \"sustainability\": \"Sustainability involves meeting present needs without compromising the ability of future generations to meet their needs.\",\n",
    "    }\n",
    "\n",
    "    topic_lower = topic.lower()\n",
    "    for key, info in research_data.items():\n",
    "        if key in topic_lower:\n",
    "            return f\"Research findings on {topic}: {info} Additional context includes recent developments and policy implications.\"\n",
    "\n",
    "    return f\"Research completed on {topic}. This is an emerging area requiring further investigation and analysis.\"\n",
    "\n",
    "\n",
    "# Analysis agent tools\n",
    "def analyze_data(research_data: str) -> str:\n",
    "    \"\"\"Analyze research data and provide insights.\"\"\"\n",
    "    if \"climate change\" in research_data.lower():\n",
    "        return \"Analysis indicates climate change requires immediate action through carbon reduction, renewable energy adoption, and international cooperation.\"\n",
    "    elif \"renewable energy\" in research_data.lower():\n",
    "        return \"Analysis shows renewable energy is becoming cost-competitive with fossil fuels and offers long-term economic and environmental benefits.\"\n",
    "    elif \"artificial intelligence\" in research_data.lower():\n",
    "        return \"Analysis reveals AI has transformative potential across industries but requires careful consideration of ethical implications and regulation.\"\n",
    "    else:\n",
    "        return \"Analysis suggests this topic has significant implications requiring strategic planning and stakeholder engagement.\"\n",
    "\n",
    "\n",
    "# Report writing agent tools\n",
    "def write_report(analysis: str, topic: str) -> str:\n",
    "    \"\"\"Write a comprehensive report based on analysis.\"\"\"\n",
    "    return f\"\"\"\n",
    "═══════════════════════════════════════\n",
    "COMPREHENSIVE RESEARCH REPORT: {topic.upper()}\n",
    "═══════════════════════════════════════\n",
    "\n",
    "EXECUTIVE SUMMARY:\n",
    "{analysis}\n",
    "\n",
    "KEY FINDINGS:\n",
    "- Evidence-based analysis indicates significant implications\n",
    "- Multiple stakeholder perspectives must be considered\n",
    "- Implementation requires coordinated approach\n",
    "- Long-term monitoring and evaluation necessary\n",
    "\n",
    "RECOMMENDATIONS:\n",
    "1. Develop comprehensive strategy framework\n",
    "2. Engage key stakeholders early in process\n",
    "3. Establish clear metrics and milestones\n",
    "4. Create feedback mechanisms for continuous improvement\n",
    "5. Allocate appropriate resources and timeline\n",
    "\n",
    "NEXT STEPS:\n",
    "- Schedule stakeholder consultations\n",
    "- Develop detailed implementation plan\n",
    "- Establish monitoring and evaluation framework\n",
    "- Begin pilot program if applicable\n",
    "\n",
    "This report provides a foundation for informed decision-decision making and strategic planning.\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "# Initialize LLM\n",
    "llm = OpenAI(model=\"gpt-4o-mini\", temperature=0)\n",
    "\n",
    "# Create individual agents using the modern API\n",
    "research_agent = FunctionAgent(\n",
    "    name=\"research_agent\",\n",
    "    description=\"This agent researches a given topic and returns key findings.\",\n",
    "    tools=[FunctionTool.from_defaults(fn=research_topic)],\n",
    "    llm=llm,\n",
    "    system_prompt=\"You are a research specialist. Use the research tool to gather comprehensive information on requested topics.\",\n",
    ")\n",
    "\n",
    "analysis_agent = FunctionAgent(\n",
    "    name=\"analysis_agent\",\n",
    "    description=\"This agent analyzes research data and provides actionable insights.\",\n",
    "    tools=[FunctionTool.from_defaults(fn=analyze_data)],\n",
    "    llm=llm,\n",
    "    system_prompt=\"You are a data analyst. Analyze research findings and provide actionable insights.\",\n",
    ")\n",
    "\n",
    "report_agent = FunctionAgent(\n",
    "    name=\"report_agent\",\n",
    "    description=\"This agent creates comprehensive, well-structured reports based on analysis.\",\n",
    "    tools=[FunctionTool.from_defaults(fn=write_report)],\n",
    "    llm=llm,\n",
    "    system_prompt=\"You are a report writer. Create comprehensive, well-structured reports based on analysis.\",\n",
    ")\n",
    "\n",
    "# Create AgentWorkflow\n",
    "multi_agent_workflow = AgentWorkflow(\n",
    "    agents=[research_agent, analysis_agent, report_agent],\n",
    "    root_agent=\"research_agent\",\n",
    ")\n",
    "\n",
    "\n",
    "async def test_agent_workflow():\n",
    "    print(\"🔍 Testing AgentWorkflow with Maxim observability...\")\n",
    "\n",
    "    query = \"\"\"I need a comprehensive report on renewable energy.\n",
    "    Please research the current state of renewable energy,\n",
    "    analyze the key findings, and create a structured report\n",
    "    with recommendations for implementation.\"\"\"\n",
    "\n",
    "    print(f\"\\n📝 Query: {query}\")\n",
    "    print(\"🔄 This will coordinate multiple agents...\")\n",
    "\n",
    "    # This will create a complex trace showing:\n",
    "    # - Multi-agent coordination\n",
    "    # - Agent handoffs and communication\n",
    "    # - Sequential tool execution\n",
    "    # - Individual agent performances\n",
    "    response = await multi_agent_workflow.run(query)\n",
    "\n",
    "    print(f\"\\n🤖 Multi-Agent Response:\\n{response}\")\n",
    "    print(\n",
    "        \"\\n✅ Check Maxim dashboard for comprehensive multi-agent workflow traces!\"\n",
    "    )\n",
    "\n",
    "\n",
    "# Run the async function\n",
    "await test_agent_workflow()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
